from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve


plt.rcParams['axes.unicode_minus'] = False
plt.rc('font', size=10)
plt.rc('font', family='SimHei')

v = 1
b = 1.7 / (2 * np.pi)
theta_0 = 0
l_head = 3.41 - 0.275 * 2
l_body = 2.2 - 0.275 * 2
l_0 = 0.275
half_width = 0.15
label_list = [0, 1, 51, 101, 151, 201, 222]


# theta与t的关系
def pfun(y, t):
    global v, b
    if t > 0:
        return v / (b * ((1 + y ** 2) ** 0.5))
    else:
        return (-1) * v / (-b * ((1 + y ** 2) ** 0.5))


def x(theta, r):
    return r * np.cos(theta)


def y(theta, r):
    return r * np.sin(theta)


# 计算速度向量的函数
def velocity_from_theta(theta, dtheta_dt, b):
    u_x = np.cos(theta) - theta * np.sin(theta)
    u_y = np.sin(theta) + theta * np.cos(theta)
    return b * dtheta_dt * np.array([u_x, u_y])


"""
求解极坐标系下圆和等距螺线的交点
"""
def find_intersections(r0, phi0, R, b, theta_guesses):
    def equations(p):
        theta = p[0]
        r = b * theta
        eq1 = r ** 2 + r0 ** 2 - 2 * r * r0 * np.cos(theta - phi0) - R ** 2
        return [eq1]

    solutions = []
    for guess in theta_guesses:
        theta_sol = fsolve(equations, guess)[0]
        r_sol = b * theta_sol
        if 0 < theta_sol < phi0 and abs(equations([theta_sol])[0]) < 1e-2:
            is_new = True
            for sol in solutions:
                if np.allclose([theta_sol], [sol[0]], atol=1e-2):
                    is_new = False
                    break
            if is_new:
                solutions.append((theta_sol, r_sol))
    solutions.sort(reverse=True)
    return solutions


# 修改后的find_next函数，现在返回位置和速度，并递归计算后续棍头
def find_next(prev_theta, prev_r, prev_pos_cart, prev_v, i=0):
    # 使用前一个棍头的极坐标作为圆心
    theta_guesses = [prev_theta - 0.0001, prev_theta - 0.001, prev_theta - 0.01, prev_theta - 0.1, prev_theta - 1]  # 初始猜测角度
    intersections = []
    if i == 0:
        intersections = find_intersections(prev_r, prev_theta, l_head, b, theta_guesses)
    else:
        intersections = find_intersections(prev_r, prev_theta, l_body, b, theta_guesses)

    results = []
    if not intersections:
        return intersections
    theta_sol, r_sol = intersections[0]
    # 当前棍头的笛卡尔坐标
    pos_cart = np.array([x(theta_sol, r_sol), y(theta_sol, r_sol)])
    # 计算向量 delta_r = current - previous
    delta_r = pos_cart - prev_pos_cart

    # 计算 u向量
    u_vec = np.array([np.cos(theta_sol) - theta_sol * np.sin(theta_sol),
                      np.sin(theta_sol) + theta_sol * np.cos(theta_sol)])

    # 计算分母和分子
    denominator = b * np.dot(u_vec, delta_r)
    numerator = np.dot(prev_v, delta_r)

    if np.abs(denominator) < 1e-6:
        # 避免除零，跳过或处理
        dtheta_dt = 0
    else:
        dtheta_dt = numerator / denominator

    # 计算当前速度
    current_v = velocity_from_theta(theta_sol, dtheta_dt, b)

    # 存储结果：位置、速度、极角、极径
    results.append({
        'theta': theta_sol,
        'r': r_sol,
        'pos_cart': pos_cart,
        'velocity': current_v,
        'index': i
    })

    if i > 222 or theta_sol > soli[999]:
        return []
    else:
        if i in label_list:
            print(f"i={i}时，位置：{x(prev_theta, prev_r), y(prev_theta, prev_r)}，速度：{(prev_v[0] ** 2 + prev_v[1] ** 2) ** 0.5}")
            print(f"b * dtheta_dt:{b * dtheta_dt}")
        # 递归计算下一个棍头，传递当前棍头的参数
        next_results = find_next(theta_sol, r_sol, pos_cart, current_v, i + 1)
        results.extend(next_results)
    return results


# 主程序
t = np.arange(0, 100, 0.1)
soli = odeint(pfun, theta_0, t)[:, 0]
t_index = 1000
theta0 = soli[t_index - 1]
r0 = b * theta0
dtheta0_dt = pfun(theta0, t[t_index - 1])

v0 = velocity_from_theta(theta0, dtheta0_dt, b)
pos0_cart = np.array([x(theta0, r0), y(theta0, r0)])

# 计算所有棍头的位置和速度
all_heads = find_next(theta0, r0, pos0_cart, v0)
# 绘制图形
plt.figure(figsize=(12, 12))
# 绘制螺线
theta_plot = np.linspace(0, 10, 1000)
r_plot = b * theta_plot
plt.plot(x(soli, b * soli), y(soli, b * soli), color='red', alpha=0.5)
plt.plot(x(soli, -b * soli), y(soli, -b * soli), color='blue', alpha=0.5)
# # 绘制第一个棍头
# plt.plot(pos0_cart[0], pos0_cart[1], 'ro', color='red')
# # 绘制其他棍头
# collision_points = []  # 观察图像可得，如果要发生碰撞，一定会发生在龙头左上方的点处，即point_4，因此记录下它的坐标
# collision_side = []  # 而可能发生碰撞的边也会是在边框的左边
# for i in range(len(all_heads)):
#     pos = all_heads[i]['pos_cart']
#     plt.plot(pos[0], pos[1], 'ro', color='red')
#     if i > 0:
#         prev_pos = all_heads[i - 1]['pos_cart']
#     else:
#         prev_pos = [pos0_cart[0], pos0_cart[1]]
#     l = np.array([pos[0] - prev_pos[0], pos[1] - prev_pos[1]])  # 由b指向a的向量
#     l_back = -1 * np.array([pos[0] - prev_pos[0], pos[1] - prev_pos[1]])  # l向量的相反向量
#     l_len = np.linalg.norm(l)
#     d = np.array([(pos[0] - prev_pos[0]) / l_len, (pos[1] - prev_pos[1]) / l_len])  # l向量平行的单位向量
#     n_left = np.array([-d[1], d[0]])  # 线段方向逆时针转得到的法向量
#     n_right = np.array([d[1], -d[0]])  # 线段方向顺时针转得到的法向量
#     n_len = np.linalg.norm(n_left)
#     point_1 = pos + l_0 * d + half_width * n_left
#     point_2 = prev_pos - l_0 * d + half_width * n_left
#     point_3 = pos + l_0 * d - half_width * n_left
#     point_4 = prev_pos - l_0 * d - half_width * n_left
#
#     # 绘制图形
#     plt.plot([point_1[0], point_3[0]], [point_1[1], point_3[1]], color="blue", alpha=0.5)
#     plt.plot([point_2[0], point_4[0]], [point_2[1], point_4[1]], color="blue", alpha=0.5)
#     plt.plot([point_3[0], point_4[0]], [point_3[1], point_4[1]], color="blue", alpha=0.5)
#     plt.plot([point_1[0], point_2[0]], [point_1[1], point_2[1]], color="black", alpha=0.5)
#     if i == 0:
#         collision_points = [point_3, point_4]
#         plt.plot(point_3[0], point_3[1], "^", color="black", markersize=4)
#         plt.plot(point_4[0], point_4[1], "^", color="black", markersize=4)
#     else:
#         collision_side.append([[point_1[0], point_1[1]], [point_2[0], point_2[1]]])
#
#
# def is_point_on_segment(point, seg_start, seg_end):
#     # 解包坐标
#     px, py = point
#     x1, y1 = seg_start
#     x2, y2 = seg_end
#
#     # 计算向量叉积（判断点是否在直线上）
#     cross_product = (px - x1) * (y2 - y1) - (py - y1) * (x2 - x1)
#     # 若叉积非零（考虑浮点精度），则点不在直线上
#     if abs(cross_product) > 0.01:
#         return False
#
#     # 检查点是否在线段的轴向包围盒内（包括端点）
#     if min(x1, x2) <= px <= max(x1, x2) and min(y1, y2) <= py <= max(y1, y2):
#         print(f"时间{t_index}时最小叉积为{cross_product}")
#         return True
#
#     return False
#
# flag = False
# # 遍历 collision_side 检测碰撞
# for collision_point in collision_points:
#     for i in range(len(collision_side)):
#         seg_start = collision_side[i][0]
#         seg_end = collision_side[i][1]
#
#         if is_point_on_segment(collision_point, seg_start, seg_end):
#             print(f"点{collision_point}与线段{seg_start} - {seg_end}发生碰撞，此时龙头的半径为{(pos0_cart[0] ** 2 + pos0_cart[1] ** 2) ** 0.5}")
#             flag = True
#             break

r_incise = 4.5
theta_incise = r_incise / b
# 切入点的笛卡尔坐标
pos_cart = np.array([x(theta_incise, r_incise), y(theta_incise, r_incise)])
plt.plot(pos_cart[0], pos_cart[1], "ro", color="blue", label="切入点")
print(f"切入点坐标{pos_cart[0], pos_cart[1]}")
plt.plot(-pos_cart[0], -pos_cart[1], "ro", color="blue", label="切出点")
plt.plot(-pos_cart[0] / 3, -pos_cart[1] / 3, "ro", color="blue", label="大小圆交点")
plt.plot([-pos_cart[0], pos_cart[0]], [-pos_cart[1], pos_cart[1]], color="black", label="弦")
string = np.array([pos_cart[0], pos_cart[1]])
string = string / np.linalg.norm(string)
r_incise_0 = 4.5 + 1e-6
theta_incise_0 = r_incise_0 / b
# 切入点附近极近一点的笛卡尔坐标
pos_cart_0 = np.array([x(theta_incise_0, r_incise_0), y(theta_incise_0, r_incise_0)])
u = np.array(pos_cart_0 - pos_cart)  # 切线向量
plt.plot([pos_cart[0] - 1e5 * u[0], pos_cart[0] + 1e5 * u[0]], [pos_cart[1] - 1e5 * u[1], pos_cart[1] + 1e5 * u[1]], color="red", label="切线1")
plt.plot([-pos_cart[0] - 1e5 * u[0], -pos_cart[0] + 1e5 * u[0]], [-pos_cart[1] - 1e5 * u[1], -pos_cart[1] + 1e5 * u[1]], color="red", label="切线2")
u = u / np.linalg.norm(u)
n_right = np.array([-u[1], u[0]])  # 逆时针转得到的法向量
plt.plot([pos_cart[0], pos_cart[0] + 4 * n_right[0]], [pos_cart[1], pos_cart[1] + 4 * n_right[1]], color="black", label="法线")
n_string = np.array([-string[1], string[0]])
plt.plot(pos_cart[0] / 3, pos_cart[1] / 3, "ro", color="blue", label="大圆弦中垂线的垂足")
plt.plot([pos_cart[0] / 3, pos_cart[0] / 3 + 2 * n_string[0]], [pos_cart[1] / 3, pos_cart[1] / 3 + 2 * n_string[1]], color="black", label="法线")
x_1 = pos_cart[0]
y_1 = pos_cart[1]
x_2 = pos_cart[0] + 2 * n_right[0]
y_2 = pos_cart[1] + 2 * n_right[1]
x_3 = pos_cart[0] / 3
y_3 = pos_cart[1] / 3
x_4 = pos_cart[0] / 3 + 2 * n_string[0]
y_4 = pos_cart[1] / 3 + 2 * n_string[1]
o_1_x = (x_2 * (y_1 - y_2) * (x_3 - x_4) - x_4 * (y_3 - y_4) * (x_1 - x_2) + (y_4 - y_2) * (x_1 - x_2) * (x_3 - x_4)) / ((y_1 - y_2) * (x_3 - x_4) - (y_3 - y_4) * (x_1 - x_2))
o_1_y = (o_1_x - x_2) * (y_1 - y_2) / (x_1 - x_2) + y_2
plt.plot(o_1_x, o_1_y, "ro", color="blue", label="大圆圆心")
o_2_x = -pos_cart[0] / 3 + (-o_1_x - pos_cart[0] / 3) / 2
o_2_y = -pos_cart[1] / 3 + (-o_1_y - pos_cart[1] / 3) / 2
plt.plot(-pos_cart[0] / 3 + (-o_1_x - pos_cart[0] / 3) / 2, -pos_cart[1] / 3 + (-o_1_y - pos_cart[1] / 3) / 2, "ro", color="blue", label="小圆圆心")
r_1 = np.linalg.norm([pos_cart[0] - o_1_x, pos_cart[1] - o_1_y])  # 大圆半径
r_2 = np.linalg.norm([-pos_cart[0] - o_2_x, -pos_cart[1] - o_2_y])  # 大圆半径
x_1 = np.arange(pos_cart[0] - 5, -pos_cart[0] / 3, 0.1)
print(x_1)
plt.plot(x_1, (r_1 ** 2 - (x_1 - o_1_x) ** 2) ** 0.5 + o_1_y, label="大圆弧")
# plt.plot(x_1, -(r_1 ** 2 - (x_1 - o_1_x) ** 2) ** 0.5 + o_1_y, label="大圆弧")
plt.xlabel('X')
plt.ylabel('Y')
plt.title('把手位置')
plt.axis('equal')
plt.grid(True)
plt.legend()
plt.show()
